import sys, os

class Diagram:
	def __init__(self,name="diagram"):
		self.out = []
		self.name = name
		self.edges = []
		self.nodes = []
		self.dot = ["digraph %s {\n" % self.name]
		return
	
	def GenDiagram(self, edges, nodes, flag=True):
		self.out = []
		if(flag):
   		    self.edges = edges
   		else:
   		    self.edges=[]
   		    
		self.nodes = nodes
		for edge in self.edges:
			out = ["\"%s\" -> \"%s\" \n" % (edge[0], edge[1])]
			self.dot += out
		
		for node in nodes:
		    out = ["\"%s\"\n" % node]
		    self.dot += out
		out = ["}\n"]
		self.dot += out
		return
	
	def WriteDiagram(self, filename=None):
		if filename:
			self.dotname = filename
		else:
			self.dotname = self.name + ".dot"
		file = open(self.dotname, 'w')
		file.writelines(self.dot)
		file.close()
		return

	def Diagram2Svg(self, filename=None):
		if filename:
			self.svgname = filename
		else:
			self.svgname = self.name + ".svg"
		os.system("dot -Tsvg %s > %s" % (self.dotname, self.svgname))

    	def GenColoredDiagram(self, edges, nodes, flag=True):
 		self.dot = ["digraph %s {\n" % self.name]
		if(flag):
   		    self.edges = edges
   		else:
   		    self.edges=[]
   		    
		self.nodes = nodes
		for edge in self.edges:
			out = ["\"%s\" -> \"%s\" [color=%s]" % (edge[0], edge[1], edge[2])]
			print out
			self.dot += out
		
		for node in nodes:
		    out = ["\"%s\"\n" % node]
		    self.dot += out
		out = ["}\n"]
		self.dot += out
		return
